import numpy as np
import math
import datetime
import pandas as pd
import json
import gc


def complex_2_float(complex_num):
    if "j" in str(complex_num):
        complex_2_str = str(complex_num)
        try:
            if "+" in complex_2_str or complex_2_str.startswith("(-"):
                complex_num_temp = complex_2_str.split("+")[0]
                if "-" in complex_num_temp:
                    complex_num_temp = complex_num_temp.split("-")[0]
                    float_num = -float(complex_num_temp.replace("(", "0"))
                else:
                    float_num = float(complex_num_temp.replace("(", "0"))

            elif "-" in complex_2_str:
                float_num = float((complex_2_str.split("-")[0]).replace("(", "0"))
        except Exception as e:
            print("error --> {} {}".format(complex_2_str, complex_num))
            print(e)
    else:
        float_num = float(complex_num)
    return float_num


def product_dict_5(xs, xlist_o, coeffs, adding_ranges=5):
    """
    xs = 根的X坐标
    ys = 根的Y坐标
    coeffs = 参数
    """
    xlist = xlist_o.copy()
    p = np.poly1d(coeffs)
    a = 0
    xsL = xs.tolist()
    tx = []
    for _xss in xsL:
        _xs = complex_2_float(_xss)
        tx.append(_xs)

    xss = max(tx)
    tx1 = tx[1]
    tx0 = tx[0]

    for i in range(adding_ranges):
        xlist.append((i + int(xss)) * 1.0)
    xlist.sort()

    newYhat = p(xlist)
    newdic = dict(zip(xlist, newYhat))

    return newdic


# def product_dict(xs, ys, coeffs):
#     """
#     xs = 根的X坐标
#     ys = 根的Y坐标
#     coeffs = 参数
#     """

#     p = np.poly1d(coeffs)
#     newX = []
#     a = 0
#     xsL = xs.tolist()
#     ysL = ys.tolist()
#     tx = []
#     ty = []
#     for _xss in xsL:
#         _xs = complex_2_float(_xss)
#         _ys = ysL[xsL.index(_xss)]
#         _ys = complex_2_float(_ys)
#         tx.append(_xs)
#         ty.append(_xs)

#     xss = max(tx)
#     tx1 = tx[1]
#     tx0 = tx[0]
#     if tx1 < 0 and tx0 < 0:
#         _tx = abs(abs(tx1) - abs(tx0))
#     elif tx1 < 0 or tx0 < 0:
#         _tx = abs(tx1) + abs(tx0)
#     else:
#         _tx = abs(tx1) - abs(tx0)

#     for i in range(int(_tx)):
#         newX.append(float(i - min(tx)))
#     newYhat = p(newX)
#     newdic = dict(zip(newX, newYhat))

#     return newdic


def polyfit(x, y, degree):
    results = {}
    coeffs = np.polyfit(x, y, degree)
    results["polynomial"] = coeffs.tolist()

    # r-squared
    p = np.poly1d(coeffs)
    Q = np.polyder(coeffs)  # 求导数
    xs = np.roots(Q)  # 求多项式函数根
    ys = np.polyval(coeffs, xs)
    # fit values, and mean

    yhat = p(x)  # or [p(z) for z in x]
    ybar = np.sum(y) / len(y)  # or sum(y)/len(y)
    ssreg = np.sum((yhat - ybar) ** 2)  # or sum([ (yihat - ybar)**2 for yihat in yhat])
    sstot = np.sum((y - ybar) ** 2)  # or sum([ (yi - ybar)**2 for yi in y])

    if sstot != 0:
        results["determination"] = ssreg / sstot  # 准确率
    else:
        results["determination"] = 0
    results["Y"] = yhat  # TODO 要删除
    results["Q"] = Q
    results["xs"] = xs
    results["ys"] = ys
    return results


def polyfitlize(df, key, gap=5, pre="5", max_incrace=0.2):
    """
    对数据进行拟合
    """

    starttime = datetime.datetime.now()
    date_list = df.date.tolist()

    total = df.Avg.dropna(axis="index")
    total = total.shape[0]

    rangeskey = {"5": "89", "15": "34", "30": "21", "60": "8", "120": "5"}
    ranges = gap
    if ranges < 10:
        ranges = ranges * 5

    prescale = 0
    loopcount = 0
    deriction = 0
    pre_deriction = 0
    Main_deriction_list = []
    deriction_list = []
    print("{} list comparing ...".format(key), end=" ")
    for _ in range(total - 1):
        scale = int((loopcount / total - 1) * 100)
        if scale % 5 == 0 and prescale != scale:
            print("*", end="")
        loopcount += 1
        prescale = scale

        deriction_list.append(_)
        compare_value_0 = df.iloc[_][key + "_ma_{}".format(rangeskey[pre])]
        compare_value_1 = df.iloc[_ + 1][key + "_ma_{}".format(rangeskey[pre])]
        if compare_value_0 >= compare_value_1:
            deriction = 1
        elif compare_value_0 <= compare_value_1:
            deriction = -1

        if pre_deriction != deriction:
            pre_deriction = deriction
            Main_deriction_list.append(deriction_list)
            deriction_list = []
            continue

    print(
        " comparing done => {}s".format((datetime.datetime.now() - starttime).seconds)
    )

    Main_deriction_list.append(deriction_list)
    temp = []
    Main_deriction_list_T = []
    num = 0
    for l in Main_deriction_list:
        num += len(l)
        temp.extend(l)
        if num > int(rangeskey[pre]):
            num = 0
            Main_deriction_list_T.append(temp)
            temp = []
    if temp != []:
        Main_deriction_list_T.append(temp)

    temp = []
    for ind in range(len(Main_deriction_list_T) - 1):
        temp = Main_deriction_list_T[ind]
        l2 = Main_deriction_list_T[ind + 1]
        lens_l1 = int(len(temp) / 2)
        lens_l2 = int(len(l2) / 2)
        temp = temp[lens_l1:]
        l2 = l2[:lens_l2]
        temp.extend(l2)
        Main_deriction_list_T.append(temp)

    if 2 < len(Main_deriction_list[-1]) < int(rangeskey[pre]):

        last_deriction_list = [
            i + int(min(Main_deriction_list[-1]) - (int(rangeskey[pre]) / 2))
            for i in range(
                int(len(Main_deriction_list[-1]) + (int(rangeskey[pre]) / 2))
            )
        ]
        if total < max(last_deriction_list):
            last_deriction_list.pop(-1)
    else:
        last_deriction_list = [total - i for i in range(int(rangeskey[pre]))]
        last_deriction_list.sort()
    Main_deriction_list_T.append(last_deriction_list)

    if gap > 40:
        Tranges = gap * 2
    elif gap > 10:
        Tranges = gap * 3
    else:
        Tranges = gap * 4
    last_p_list_Main = []
    for n in range(20):
        last_p_list = [total - i - n for i in range(int(Tranges))]
        last_p_list.sort()
        Main_deriction_list_T.append(last_p_list)
        last_p_list_Main.append(last_p_list)

    Main_deriction_list_T.sort()

    l = {}
    YDic = {}
    determination_dic = {}
    compare_dic_con = {}
    compare_dic_con_d = {}
    XYdic = {}
    print("{} polyfitlize start".format(key), end=" ")
    for deriction_list in Main_deriction_list_T:
        scale = int((loopcount / len(Main_deriction_list_T)) * 100)
        if scale % 5 == 0 and prescale != scale:
            print("*", end="")
        loopcount += 1
        prescale = scale
        y = df.iloc[deriction_list][key]
        y.dropna(axis=0, inplace=True)
        while np.nan in y:
            y.remove(np.nan)

        dic = {}
        x = []
        a = 0
        for i in y:
            x.append(float(a))
            a += 1

        results = polyfit(x, y, 3)
        dic["determination_T"] = results["determination"]
        a = 0
        coeffsL = results["polynomial"]

        for i in coeffsL:
            dic["polynomial_X_{}".format(a)] = coeffsL[a]
            a += 1

        a = 0
        Ql = results["Q"]
        for i in Ql:
            dic["Q_{}".format(a)] = Ql[a]
            a += 1

        a = 0
        xsL = results["xs"].tolist()
        ysL = results["ys"].tolist()
        deriction_list_min = deriction_list[0]
        temp_dic = {}
        for xs in xsL:
            _xs = round(complex_2_float(xs) + deriction_list_min)
            if 0 < _xs < df.shape[0]:
                ys = ysL[xsL.index(xs)]
                ys = complex_2_float(ys)
                temp_dic.update({df.iloc[int(_xs)].date: ys})
                if int(_xs) > deriction_list_min:
                    day_gap = int(
                        df.date.iloc[deriction_list_min : int(_xs)].shape[0] / ranges
                    )
                elif int(_xs) == deriction_list_min:
                    day_gap = 1
                else:
                    day_gap = int(
                        df.date.iloc[int(_xs) : deriction_list_min].shape[0] / ranges
                    )
                if (
                    df.iloc[deriction_list_min].Avg
                    * ((1 - max_incrace) ** abs(day_gap))
                    <= ys
                    <= df.iloc[deriction_list_min].Avg
                    * ((1 + max_incrace) ** abs(day_gap))
                    and day_gap < (240 / ranges) * 5
                ):
                    XYdic[df.iloc[deriction_list_min].date] = temp_dic

        if deriction_list == last_deriction_list:
            last_deriction_list_date = df.iloc[min(last_deriction_list)].date
            YDic[last_deriction_list_date] = product_dict_5(
                xs=results["xs"],
                xlist_o=x,
                coeffs=results["polynomial"],
                adding_ranges=(
                    max(last_deriction_list) - df.shape[0]
                    if max(last_deriction_list) - df.shape[0] > 5
                    else ranges
                ),
            )
        else:
            T = dict(zip(x, results["Y"]))
            YDic[df.iloc[deriction_list_min].date] = T

        if (
            deriction_list in last_p_list_Main
            and 0.70 < results["determination"] < 1
        ):

            compare_dic_con[df.iloc[deriction_list_min].date] = product_dict_5(
                xs=results["xs"],
                xlist_o=x,
                coeffs=results["polynomial"],
                adding_ranges=int(gap / 2),
            )
            compare_dic_con_d[df.iloc[deriction_list_min].date] = results[
                "determination"
            ]

        l[df.iloc[deriction_list[0]]["date"]] = dic
        determination_dic[df.iloc[deriction_list_min].date] = results["determination"]
    print(" polyfitlize => {}s".format((datetime.datetime.now() - starttime).seconds))

    print("Y_prodict   start", end=" ")
    new_dict = {}
    date_list = df.date.tolist()
    predic = {}
    loopcount = 0
    YDL = list(YDic.keys())
    YDL.sort()
    for date in YDL:
        scale = int((loopcount / len(YDic.keys())) * 100)
        if scale % 5 == 0 and prescale != scale:
            print("*", end="")
        loopcount += 1
        prescale = scale

        index = date_list.index(date)
        value = YDic[date]
        ls_value = l[date]
        for ind in value:
            if 0 <= index + int(ind) < len(date_list):
                if date_list[index + int(ind)] not in new_dict:
                    new_dict[date_list[index + int(ind)]] = {
                        "Y_prodict": round(value[int(ind)], 2),
                        "determination": determination_dic[date],
                        "prodict_date": date,
                        "polynomial_X_3": ls_value["polynomial_X_3"],
                        "polynomial_X_2": ls_value["polynomial_X_2"],
                        "polynomial_X_1": ls_value["polynomial_X_1"],
                        "Q_2": ls_value["Q_2"],
                        "Q_1": ls_value["Q_1"],
                        "Q_0": ls_value["Q_0"],
                    }
                elif (
                    new_dict[date_list[index + int(ind)]]["determination"]
                    < determination_dic[date]
                ):

                    new_dict[date_list[index + int(ind)]] = {
                        "Y_prodict": round(value[int(ind)], 2),
                        "determination": determination_dic[date],
                        "prodict_date": date,
                        "polynomial_X_3": ls_value["polynomial_X_3"],
                        "polynomial_X_2": ls_value["polynomial_X_2"],
                        "polynomial_X_1": ls_value["polynomial_X_1"],
                        "Q_2": ls_value["Q_2"],
                        "Q_1": ls_value["Q_1"],
                        "Q_0": ls_value["Q_0"],
                    }
    value = YDic[last_deriction_list_date]
    index = date_list.index(last_deriction_list_date)
    for ind in value:
        if 0 <= index + int(ind) < len(date_list):
            if int(ind) < 2:
                new_dict[date_list[index + int(ind)]] = {
                    "Y_prodict": round(value[int(ind)], 2),
                    "determination": determination_dic[last_deriction_list_date],
                    "prodict_date": last_deriction_list_date,
                    "polynomial_X_3": ls_value["polynomial_X_3"],
                    "polynomial_X_2": ls_value["polynomial_X_2"],
                    "polynomial_X_1": ls_value["polynomial_X_1"],
                    "Q_2": ls_value["Q_2"],
                    "Q_1": ls_value["Q_1"],
                    "Q_0": ls_value["Q_0"],
                }
            else:
                new_dict[date_list[index + int(ind)]] = {
                    "Y_prodict": round(value[int(ind)], 2),
                    "prodict_date": last_deriction_list_date,
                }

    compare_dic = {}
    for p_date in compare_dic_con:
        value = compare_dic_con[p_date]
        index = date_list.index(p_date)
        for ind in value:
            if 0 <= index + int(ind) < len(date_list):
                compare_dic[date_list[index + int(ind)]] = {
                    "Y_prodict_compare": round(value[int(ind)], 2),
                    "determination_compare": compare_dic_con_d[p_date],
                    "prodict_date_compare": p_date,
                }

    if compare_dic == {}:
        compare_dic[date_list[index + int(ind)]] = {
            "Y_prodict_compare": np.nan,
            "determination": np.nan,
            "prodict_date_compare": np.nan,
        }

    print("Y_prodict  => {}s".format((datetime.datetime.now() - starttime).seconds))

    print(date_list[-1], len(date_list))
    _dfY = pd.DataFrame(new_dict).T.reset_index()
    _df_compare = pd.DataFrame(compare_dic).T.reset_index()
    _dfY.rename(columns={"index": "date"}, inplace=True)
    _df_compare.rename(columns={"index": "date"}, inplace=True)
    df = pd.merge(df, _dfY, on="date", how="left")
    df = pd.merge(df, _df_compare, on="date", how="left")
    print(df.columns.tolist())
    # for i in
    # df.dropna(
    #     axis="index",
    #     how="all",
    #     subset=["Avg", "Y_prodict", "Y_prodict_compare"],
    #     inplace=True,
    # )

    new_l = []

    for XY in XYdic:
        dic = {}
        dic["date"] = XY
        if XYdic[XY] != {}:
            a = 0
            for _ in XYdic[XY]:
                dic["xss_{}".format(a)] = _
                dic["ys_{}".format(a)] = XYdic[XY][_]
                a += 1
            new_l.append(dic)

    XYdic = pd.DataFrame(new_l)
    # XYdic.rename(columns={"index": "date"}, inplace=True)
    df = pd.merge(df, XYdic, on="date", how="left")

    df.to_csv("Main.csv")
    return df